import socket
import struct
import time

class ModbusClient:
    def __init__(self, host, port=4001):
        self.host = host
        self.port = port
        self.sock = None
        self.transaction_id = 0
        
    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, self.port))
        
    def close(self):
        if self.sock:
            self.sock.close()
            
    def _get_transaction_id(self):
        """获取自增的事务ID"""
        self.transaction_id = (self.transaction_id + 1) & 0xFFFF
        return self.transaction_id
        
    def write_firmware_block(self, block_num, data):
        """写入固件块
        
        Args:
            block_num: 块号(0-59)
            data: 1KB的数据
            
        Returns:
            True表示成功，False表示失败
        """
        tx_data = struct.pack('>HHHBB',
            self._get_transaction_id(),  # 事务ID
            0,                          # 协议ID
            len(data) + 3,             # 长度 (数据长度 + 功能码长度 + 块号长度)
            1,                         # 单元ID
            0x50                       # 功能码 - 写入固件
        )
        
        # 添加块号和数据
        tx_data += bytes([block_num]) + data
        
        # 发送请求
        self.sock.send(tx_data)
        
        # 接收响应
        response = self.sock.recv(1024)
        if len(response) < 9:
            raise Exception("Response too short")
            
        # 检查状态
        return response[8] == 0
        
    def verify_firmware(self, crc):
        """验证固件
        
        Args:
            crc: 32位CRC值
            
        Returns:
            True表示验证成功，False表示失败
        """
        tx_data = struct.pack('>HHHBBI',
            self._get_transaction_id(),  # 事务ID
            0,                          # 协议ID
            6,                         # 长度 (功能码 + CRC)
            1,                         # 单元ID
            0x51,                      # 功能码 - 验证固件
            crc                        # CRC值
        )
        
        # 发送请求
        self.sock.send(tx_data)
        
        # 接收响应
        response = self.sock.recv(1024)
        if len(response) < 9:
            raise Exception("Response too short")
            
        # 检查状态
        return response[8] == 0
        
    def activate_firmware(self):
        """激活固件更新
        
        Returns:
            True表示成功，False表示失败
        """
        tx_data = struct.pack('>HHHBB',
            self._get_transaction_id(),  # 事务ID
            0,                          # 协议ID
            2,                         # 长度 (仅功能码)
            1,                         # 单元ID
            0x52                       # 功能码 - 激活固件
        )
        
        # 发送请求
        self.sock.send(tx_data)
        
        # 接收响应
        response = self.sock.recv(1024)
        if len(response) < 9:
            raise Exception("Response too short")
            
        # 检查状态
        return response[8] == 0
